package com.example.common.utils;

import com.google.gson.Gson;
import lombok.AllArgsConstructor;
import lombok.Data;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * Rsa: 使用PKCS#8标准
 */
@SuppressWarnings ("unused")
public abstract class RsaUtils {
    private RsaUtils() {}

    /**
     * 加密算法
     */
    private static final String ALGORITHM = "RSA";

    /**
     * 加密模式
     */
    private static final String CIPHER_MODE = "RSA/ECB/PKCS1Padding";
    /**
     * 签名算法
     */
    private static final String SIGN_ALGORITHM = "SHA256withRSA";

    /**
     * 公私钥类
     */
    @Data
    @AllArgsConstructor
    public static class Keypair {
        private final String privateKey;
        private final String publicKey;
    }

    /**
     * 生成PKCS#1标准格式的公私钥
     *
     * @return 公钥与私钥, 均为base64编码后的结果
     */
    public static Keypair generate() throws Exception {
        // 获取生成器并生成键值对对象
        KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM);
        generator.initialize(4096, new SecureRandom());
        KeyPair pair = generator.generateKeyPair();
        // 获取pkcs#8标准的私钥
        PrivateKey pri = pair.getPrivate();
        String privateKey = Base64.getEncoder().encodeToString(pri.getEncoded());
        // 获取pkcs#8标准的公钥
        PublicKey pub = pair.getPublic();
        String publicKey = Base64.getEncoder().encodeToString(pub.getEncoded());
        return new Keypair(privateKey, publicKey);
    }

    /**
     * 使用公钥对数据进行加密
     * 注意该方法会使用gson尝试将传入数据转为json字符串随后对json字符串进行加密
     *
     * @param data      数据明文
     * @param publicKey base64编码的公钥
     * @return base64编码的密文
     */
    public static String encrypt(Object data, String publicKey) throws Exception {
        // 将传入数据转为json字符串
        String json = new Gson().toJson(data);
        // 对json字符串进行加密
        byte[] encrypt = encrypt(json.getBytes(StandardCharsets.UTF_8), publicKey);
        // 返回base64编码后的密文
        return Base64.getEncoder().encodeToString(encrypt);
    }

    /**
     * 使用公钥对数据进行加密
     *
     * @param data      数据明文
     * @param publicKey base64编码的公钥
     * @return 密文
     */
    public static byte[] encrypt(byte[] data, String publicKey) throws Exception {
        // 获取公钥实例对象
        PublicKey pub = getPublicKey(publicKey);
        // 获取密码器实例
        Cipher cipher = Cipher.getInstance(CIPHER_MODE);
        // 初始化密码器: 指定模式为加密并设置公钥
        cipher.init(Cipher.ENCRYPT_MODE, pub);
        // 对数据进行加密
        return cipher.doFinal(data);
    }

    /**
     * 使用私钥对数据进行解密
     * 该方法首先将base64编码的密文进行解码, 在解密后会使用gson将json字符串转为目标数据类型
     *
     * @param encrypt    base64编码的数据密文
     * @param privateKey base64编码的私钥
     * @param target     目标数据类型
     * @return base64编码后的密文
     */
    public static <T> T decrypt(String encrypt, String privateKey, Class<T> target) throws Exception {
        // 解密数据
        byte[] decrypt = decrypt(Base64.getDecoder().decode(encrypt), privateKey);
        // 将字节数据转为json字符串
        String json = new String(decrypt, StandardCharsets.UTF_8);
        // 将json字符串转为目标类型
        return new Gson().fromJson(json, target);
    }

    /**
     * 使用私钥对数据进行解密
     *
     * @param encrypt    数据密文
     * @param privateKey base64编码的私钥
     * @return 明文
     */
    public static byte[] decrypt(byte[] encrypt, String privateKey) throws Exception {
        // 获取私钥实例对象
        PrivateKey pri = getPrivateKey(privateKey);
        // 获取密码器实例
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        // 初始化密码器, 指定模式为解密并设置私钥
        cipher.init(Cipher.DECRYPT_MODE, pri);
        // 对数据进行解密
        return cipher.doFinal(encrypt);
    }

    /**
     * 使用私钥通过SHA256算法对数据进行签名
     * 注意该方法会使用gson尝试将传入数据转为json字符串随后对json字符串进行签名
     *
     * @param data       数据明文
     * @param privateKey base64编码的私钥
     * @return base64编码的签名
     */
    public static String sign(Object data, String privateKey) throws Exception {
        // 将传入数据转为json字符串
        String json = new Gson().toJson(data);
        // 对json字符串进行签名
        byte[] sign = sign(json.getBytes(StandardCharsets.UTF_8), privateKey);
        // 返回base64编码后的签名
        return Base64.getEncoder().encodeToString(sign);
    }

    /**
     * 使用私钥通过SHA256算法对数据进行签名
     *
     * @param data       数据明文
     * @param privateKey base64编码的私钥
     * @return 签名
     */
    public static byte[] sign(byte[] data, String privateKey) throws Exception {
        // 获取私钥实例对象
        PrivateKey pri = getPrivateKey(privateKey);
        // 获取签章实例
        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        // 使用私钥初始化签章进行签名
        signature.initSign(pri);
        // 执行签名
        signature.update(data);
        return signature.sign();
    }

    /**
     * 使用公钥通过SHA256算法对数据签名进行验签
     * 注意:
     * 该方法会使用gson尝试将传入数据转为json字符串
     * 将签名使用base64来解码随后进行验签
     *
     * @param data      数据明文
     * @param publicKey base64编码的公钥
     * @return base64编码的密文
     */
    public static boolean verify(Object data, String sign, String publicKey) throws Exception {
        // 将传入数据转为json字符串
        String json = new Gson().toJson(data);
        // 将签名进行解码
        byte[] s = Base64.getDecoder().decode(sign);
        // 验签
        return verify(json.getBytes(StandardCharsets.UTF_8), s, publicKey);
    }

    /**
     * 使用公钥通过SHA256算法对数据签名进行验签
     *
     * @param data      数据明文
     * @param sign      签名
     * @param publicKey base64编码的公钥
     * @return 验签是否通过(数据的签名是否来自公钥对应的私钥)
     */
    public static boolean verify(byte[] data, byte[] sign, String publicKey) throws Exception {
        // 获取公钥实例对象
        PublicKey pub = getPublicKey(publicKey);
        // 获取签章实例
        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        // 使用公钥初始化签章进行验签
        signature.initVerify(pub);
        // 执行验签
        signature.update(data);
        return signature.verify(sign);
    }

    /**
     * 将base64编码的私钥解码并生成私钥实例对象
     *
     * @param privateKey base64编码的私钥
     * @return 私钥实例对象
     */
    private static PrivateKey getPrivateKey(String privateKey) throws Exception {
        // 将base64编码的私钥解码
        byte[] pkb = Base64.getDecoder().decode(privateKey);
        // 通过私钥字节数据获取私钥实例对象
        return KeyFactory.getInstance(ALGORITHM).generatePrivate(new PKCS8EncodedKeySpec(pkb));
    }

    /**
     * 将base64编码的公钥解码并生成公钥实例对象
     *
     * @param publicKey base64编码的公钥
     * @return 私钥实例对象
     */
    private static PublicKey getPublicKey(String publicKey) throws Exception {
        // 将base64编码的公钥解码
        byte[] pkb = Base64.getDecoder().decode(publicKey);
        // 获取公钥实例对象
        return KeyFactory.getInstance(ALGORITHM).generatePublic(new X509EncodedKeySpec(pkb));
    }

    public static void main(String[] args) throws Exception {
        long startTime = System.currentTimeMillis();
        String data = "Hello World";
        System.out.println("初始数据: " + data);

        Keypair pair = generate();
        System.out.println("私钥: " + pair.privateKey);
        System.out.println("公钥: " + pair.publicKey);

        String encrypt = encrypt(data, pair.publicKey);
        System.out.println("加密后的数据: " + encrypt);

        String decrypt = decrypt(encrypt, pair.privateKey, String.class);
        System.out.println("解密后的数据: " + decrypt);

        String sign = sign(data, pair.privateKey);
        System.out.println("签名后的数据: " + sign);

        System.out.println("验签结果: " + verify(data, sign, pair.publicKey));

        System.out.println("消耗的时间: " + (System.currentTimeMillis() - startTime) + "ms");
    }

}
